# app/rag/embedding/huggingface.py

from .base import EmbeddingProvider
from langchain_community.embeddings import HuggingFaceEmbeddings
from typing import List
import torch


class HuggingFaceEmbedding(EmbeddingProvider):
    def __init__(self, model_name: str = "BAAI/bge-small-zh-v1.5", device: str = "cpu"):
        self.embedding = HuggingFaceEmbeddings(
            model_name=model_name,
            model_kwargs={"device": device}
        )

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.embedding.embed_documents(texts)

    def embed_query(self, text: str) -> List[float]:
        return self.embedding.embed_query(text)